module Zeno.Core where

import Prelude ()
import StdImports
import Data.IntMap ( IntMap )
import qualified Data.IntMap as IMap
import qualified Data.Map as Map
import qualified Data.Set as Set

import Zeno.Clause
import Zeno.Type
import Zeno.Term
import Zeno.Function
import Zeno.ProofLogic

debug = False

type ZTerm = Term ZVar
type ZType = Type ZTypeVar
type ZClause = Clause ZVar
type ZEquality = Equality ZVar
type ZQuantifiedClause = Quantified Clause ZVar
type ZTermSubstitution = TermSubstitution ZVar
type ZTypeSubstitution = TypeSubstitution ZTypeVar
type ZProofStep = ProofStep ZVar
type ZProof = Proof ZVar
type ZFunction = Function ZVar
type ZExpression = Expression ZVar
type ZWhen = When ZVar
type VariableSource = [Id]
type CriticalTerm = (ZTerm, VariableSource)

data ZVar
  = ZVar        { zVarId :: !Id,
                  zVarName :: !(Maybe String),
                  zVarType :: !ZType,
                  zVarClass :: !ZVarClass }
                                    
instance Eq ZVar where
  var1 == var2 = zVarId var1 == zVarId var2
  
instance Ord ZVar where
  var1 `compare` var2 = 
    zVarId var1 `compare` zVarId var2
    
instance Show ZVar where
  show var = 
    let name = 
          case zVarName var of
            Just name -> name
            Nothing -> show (zVarId var)
        sources = 
          case zVarClass var of
            ZVariable srs@(_:_) -> ("{" ++)
                                 $ (++ "}")
                                 $ intercalate ", " 
                                 $ map show srs
            _ -> ""
    in name -- ++ sources

data ZVarClass
  = ZVariable     { zVarSources :: ![VariableSource] }
  | ZFunction     { zVarFunction :: !ZFunction }
  | ZConstructor
  deriving ( Eq )

data ZTypeVar
  = ZDataType   { zDataTypeId :: !Id,
                  zDataTypeName :: !String,
                  zDataTypeArgs :: ![ZTypeVar],
                  zDataTypeCons :: ![ZVar] }
                  
  | ZPolyType   { zPolyTypeId :: !Id }

instance Show ZTypeVar where
  show (ZDataType _ name _ _) = name
  show (ZPolyType id) = show id
  
instance Eq ZTypeVar where
  (ZDataType id1 _ _ _) == (ZDataType id2 _ _ _) = id1 == id2
  (ZPolyType id1) == (ZPolyType id2) = id1 == id2
  _ == _ = False
  
instance Ord ZTypeVar where
  (ZPolyType {}) `compare` (ZDataType {}) = LT
  (ZDataType {}) `compare` (ZPolyType {}) = GT
  (ZPolyType id1) `compare` (ZPolyType id2) = id1 `compare` id2
  (ZDataType id1 _ _ _) `compare` (ZDataType id2 _ _ _) = id1 `compare` id2
  
data Zeno
  = Zeno        { zFunctions :: !(Map String ZFunction),
                  zDataTypes :: !(Map String ZTypeVar),
                  zCounter :: !Id }

instance IdCounter Zeno where
  newId zeno = 
    let id = nextId (zCounter zeno)
    in (id, zeno { zCounter = id })
                  
emptyZeno 
  = Zeno        { zFunctions = mempty,
                  zDataTypes = mempty,
                  zCounter = firstId }
                  
type ZenoState = State Zeno
type ZenoReader = Reader Zeno
                                    
polyTypeVar :: ZType -> Bool
polyTypeVar (TypeVar (ZPolyType _)) = True
polyTypeVar _  = False

dataTypeVar :: ZType -> Bool
dataTypeVar (TypeVar (ZDataType {})) = True
dataTypeVar _ = False

isFunctionClass :: ZVarClass -> Bool
isFunctionClass (ZFunction {}) = True
isFunctionClass _ = False

isVariableClass :: ZVarClass -> Bool
isVariableClass (ZVariable {}) = True
isVariableClass _ = False

isConstructorClass :: ZVarClass -> Bool
isConstructorClass ZConstructor = True
isConstructorClass _ = False

functionVar :: ZVar -> Bool
functionVar = isFunctionClass . zVarClass

constructorVar :: ZVar -> Bool
constructorVar = isConstructorClass . zVarClass

functionTerm :: ZTerm -> Bool
functionTerm = functionVar . termFunction

groundTerm :: ZTerm -> Bool
groundTerm = all ((== ZConstructor) . zVarClass)

unequalTerms :: ZTerm -> ZTerm -> Bool
unequalTerms t1 t2 = 
  let ((Var func1) : args1) = flattenTerm t1
      ((Var func2) : args2) = flattenTerm t2
  in if not (constructorVar func1 && constructorVar func2)
    then False
    else if func1 /= func2
      then True
      else or (zipWith unequalTerms args1 args2)
      
unequal :: ZEquality -> Bool
unequal (t1 `Equals` t2) = unequalTerms t1 t2

equal :: ZEquality -> Bool 
equal (t1 `Equals` t2) = equalTerms t1 t2

mergeCriticalTerms :: [CriticalTerm] -> [CriticalTerm]
mergeCriticalTerms cterms = 
  let clustered = clusterBy ((==) `on` fst) cterms
  in  concatMap (filterEach requiredSource) clustered
  where
    requiredSource :: (CriticalTerm, [CriticalTerm]) -> Bool
    requiredSource ((_, src), other_terms) =
      let other_srcs = map snd other_terms
      in  not 
        $ any (flip orderedSupersetOf src) 
        $ other_srcs

recursiveType :: ZType -> Bool
recursiveType ztype = 
  let (type_func : _) = flattenAppType ztype
  in  case type_func of
        TypeVar tvar -> recursiveTypeVar tvar
        _ -> False
 
recursiveTypeVar :: ZTypeVar -> Bool
recursiveTypeVar (ZPolyType {}) = False
recursiveTypeVar dt@(ZDataType {}) = 
  any (elem dt) $ map toArgTypes $ zDataTypeCons dt
  where 
    toArgTypes = concatMap toList . butlast . flattenFunType . zVarType

variables :: Foldable f => f ZVar -> [ZVar]
variables = nub . filter (isVariableClass . zVarClass) . toList

constructors :: Foldable f => f ZVar -> [ZVar]
constructors = nub . filter (isConstructorClass . zVarClass) . toList

variableSources :: ZVar -> [VariableSource]
variableSources var = 
  case zVarClass var of
    ZVariable srs -> srs
    _ -> []
    
removeSources :: ZVar -> ZVar
removeSources var = 
  case zVarClass var of
    ZVariable _ -> var 
      { zVarClass = ZVariable [] } 
    _ -> var
    
addSources :: [VariableSource] -> ZVar -> ZVar
addSources new_srs var = 
  case zVarClass var of
    ZVariable existing_srs -> 
      let all_srs = nub (new_srs ++ existing_srs)
      in var { zVarClass = ZVariable all_srs }
    _ -> 
      var
      
usableCriticalTerm :: CriticalTerm -> Bool
usableCriticalTerm (term, src) 
   = all (not . any (orderedSupersetOf src) . variableSources) term
  && not (null src)
          
substituteTakingSources :: WithinTraversable ZTerm t =>
    ZTerm -> ZTerm -> t -> t
substituteTakingSources from to = mapWithin substitute
  where
    substitute term 
      | term == from = 
          let sources = concatMap variableSources
                      $ toList term
          in fmap (addSources sources) to
      | otherwise = term   

variableTerm :: ZTerm -> Bool
variableTerm (Var var) = isVariableClass (zVarClass var)
variableTerm _ = False

defaultVariable :: ZVarClass
defaultVariable = ZVariable []

constructorTerm :: ZTerm -> Bool
constructorTerm term = 
  zVarClass (termFunction term) == ZConstructor

freshVariable :: ZVar -> ZenoState ZVar
freshVariable (ZVar id _ typ cls) = do
  new_id <- newIdS
  return $ ZVar new_id Nothing typ cls

lookupFunction :: String -> ZenoReader ZFunction
lookupFunction name = do
  funcs <- zFunctions <$> ask
  case Map.lookup name funcs of
    Just fun -> return fun
    Nothing -> error $ "Couldn't find function " ++ name

instantiateConstructors :: forall t . IdCounter t =>
    CriticalTerm -> Maybe (State t [(ZTerm, [ZVar])])
instantiateConstructors (term, source) =
  if not (dataTypeVar ty_func) 
     then Nothing
     else Just (mapM instantiate cons)
  where
    ty_func : ty_args = flattenAppType (termType term)
    TypeVar dtype@(ZDataType _ _ dt_args cons) = ty_func
    bindings = Set.fromList $ (map TypeVar dt_args) `zip` ty_args
    
    instantiate :: ZVar -> State t (ZTerm, [ZVar])
    instantiate var = do
      let var_type = substituteWithin bindings (zVarType var)
          var_types = flattenFunType var_type
          (arg_types, [res_type]) = splitAt (length var_types - 1) var_types
      new_ids <- replicateM (length arg_types) newIdS
      let var_class = ZVariable [source]
          makeVar id typ = ZVar id Nothing typ var_class
          args = zipWith ($) (map makeVar new_ids) arg_types
          term = unflattenTerm (map Var (var : args))
          vars = filter ((== res_type) . zVarType) args
      return (term, vars)

validClause :: ZClause -> Bool
validClause (Clause head body) = all validBodyEq body
  where validBodyEq :: ZEquality -> Bool
        validBodyEq (Equals left_term right_term) = 
          functionTerm left_term && not (functionTerm right_term)
          
termType :: ZTerm -> ZType
termType term =
  let ((Var fun) : args) = flattenTerm term
      fun_types = flattenFunType (zVarType fun)
      res_types = drop (length args) fun_types
  in unflattenFunType res_types
 
destructibleType :: ZType -> Bool
destructibleType = dataTypeVar . head . flattenAppType

destructibleTerm :: ZTerm -> Bool
destructibleTerm term 
   = destructibleType (termType term)
  && not (isConstructorClass (zVarClass (termFunction term)))

destructibleVar :: ZVar -> Bool
destructibleVar var 
   = isVariableClass (zVarClass var)
  && destructibleType (zVarType var)
  
addDataType :: MonadState Zeno m => ZTypeVar -> m ()
--addDataType (ZPolyType _) = return ()
addDataType dtype@(ZDataType _ name _ _) = modify $ \z -> z 
  { zDataTypes = Map.insert name dtype (zDataTypes z) }

addFunction :: ZFunction -> ZenoState ()
addFunction def = do
  let Just name = zVarName (funcName def)
  modify $ \z -> z 
    { zFunctions = Map.insert name def (zFunctions z) }

clausifyFunction :: ZFunction -> [ZClause]
clausifyFunction (Function _ _ (Value (Var (ZVar _ _ _ (ZFunction func))))) =
  clausifyFunction func
clausifyFunction (Function _ _ expr) =
  map normaliseClause (clausifyExpression expr)
  where
    clausifyExpression :: ZExpression -> [ZClause]
    clausifyExpression (Value (App (App (Var _) x) y)) = 
      [ Clause (x `Equals` y) [] ]
    clausifyExpression (Value _) = []
    clausifyExpression (CaseOf _ ct ws) = concatMap clausifyWhen ws
      where
        clausifyWhen (When wt expr) = 
          map (addCondition (ct `Equals` wt)) (clausifyExpression expr)
          
    normaliseClause :: ZClause -> ZClause
    normaliseClause cls@(Clause eq conds) =
      case find (variableTerm . equalityLeft) conds of
        Just (from `Equals` to) ->
          let replace = replaceWithin from to
              conds' = map replace conds
              conds'' = filter (not . uncurry equalTerms . equalityToPair) conds'
          in normaliseClause (Clause eq conds'')
        Nothing -> 
          let (conds', any) = runWriter $ concatMapM factoriseEquality conds
          in  if getAny any 
                then normaliseClause (Clause eq conds')
                else cls
      where
        factoriseEquality :: ZEquality -> Writer Any [ZEquality]
        factoriseEquality eq@(t1 `Equals` t2) = 
          let ((Var func1) : args1) = flattenTerm t1
              ((Var func2) : args2) = flattenTerm t2
          in  if func1 == func2 && zVarClass func1 == ZConstructor
                then do
                  tell (Any True)
                  let pairs = args1 `zip` args2
                  return $ fmap pairToEquality pairs
                else return [eq]
  
instance Show Zeno where
  show zeno = "\n\n" ++ datatypes_s ++ 
              "\n\n" ++ functions_s
    where
      datatypes_s = intercalate "\n\n" $ 
        map showFullZTypeVar $ Map.elems $ zDataTypes zeno
        
      functions_s = intercalate "\n\n" $ 
        map show $ Map.elems $ zFunctions zeno
                
showTypedZTerm :: ZTerm -> String
showTypedZTerm term = 
  show term ++ " :: " ++ show (termType term)
  
showFullZTypeVar :: ZTypeVar -> String
showFullZTypeVar ptype@(ZPolyType {}) = show ptype
showFullZTypeVar dtype@(ZDataType _ name args cons) =
  rec_s ++ "type " ++ name ++ " " 
    ++ args_s ++ "\n  " ++ cons_s
  where
    args_s = intercalate " " $ map show args
    cons_s = intercalate "\n  " $ map (showTypedZTerm . Var) cons
    rec_s = if recursiveTypeVar dtype then "rec " else ""
    
showZProof :: ZProof -> String
showZProof proof =
  show proof' ++ "\n\nProven: " ++ inds_s
  where
    ([proof'], inds) = runWriter $ findInducts True proof
    inds_s = intercalate "\n        " 
      $ map show inds
  
    findInducts :: Bool -> ZProof -> Writer [ZClause] [ZProof]
    findInducts _ Unknown = return [Unknown]
    findInducts _ ds@(Disproven {}) = return [ds] 
    findInducts ind pr@(Proven step@(ProofStep _ Nothing []) ps) = do
      ps' <- concatMapM (findInducts True) ps
      return [Proven step ps']
    findInducts ind (Proven (ProofStep "fct" _ []) ps) =
      concatMapM (findInducts ind) ps
    findInducts ind (Proven step@(ProofStep name (Just cls) _) ps) = do
      let ind_step = name == "ind" || name == "cse"
          ind' = not ind_step && 
            ("hyp" `isPrefixOf` name || 
             "gen" `isPrefixOf` name || ind)
      ps' <- concatMapM (findInducts ind') ps
      when (ind_step && ind) (tell [cls])
      if ind_step
        then return ps'
        else return [Proven step ps']
        
-- Debug section
  
traceZ :: String -> a -> a
traceZ 
  | debug = trace
  | otherwise = flip const

